#%%
import numpy as np
import h5py
import tqdm
import pickle
import os.path as osp

# %%
data_h5_file='/DATA/taoxianming/rat/data/Mix_analysis/SexAgeDay55andzzcWTinAUT_MMFF/result32/semiseq2seq_iter0/datasets/data-iter3.h5'

## --------- read data
with h5py.File(data_h5_file, 'r') as f:
    assert '0' in f and '66666' not in f and 'label' in f
    label = np.array(f['label'])
    nsample = len(label)
    feat_raw = np.array([f[str(i)] for i in tqdm.trange(nsample)])


assert 0 in label
ind_has_anno = label>0
label_not0 = label[ind_has_anno]
feat_raw_not0 = feat_raw[ind_has_anno]
nsample_not0 = len(label_not0)

## ------- create train test dataset
isamples = np.arange(nsample_not0)
np.random.seed(10)
np.random.shuffle(isamples)
ntrain = int(nsample*0.8)
isamples_train = np.sort(isamples[:ntrain])
isamples_test = np.sort(isamples[ntrain:])
label_train = label_not0[isamples_train]
label_test = label_not0[isamples_test]
feat_raw_train = feat_raw_not0[isamples_train]
feat_raw_test = feat_raw_not0[isamples_test]

## ------- write to h5 file
data_h5_dir = osp.dirname(data_h5_file)
data_h5_train_file = osp.join(data_h5_dir, 'data-iter3-train.h5')
data_h5_test_file = osp.join(data_h5_dir, 'data-iter3-test.h5')

with h5py.File(data_h5_train_file, 'w') as f:
    f.create_dataset('label', data=label_train)
    for i in tqdm.trange(len(feat_raw_train)):
        f.create_dataset(str(i), data=feat_raw_train[i])

        
with h5py.File(data_h5_test_file, 'w') as f:
    f.create_dataset('label', data=label_test)
    for i in tqdm.trange(len(feat_raw_test)):
        f.create_dataset(str(i), data=feat_raw_test[i])

meta_pkl = osp.join(data_h5_dir, 'data-iter3-nantraintest.pkl')
index = np.zeros_like(label) + np.nan
x = np.zeros([ind_has_anno.sum()])
x[isamples_train]==0
x[isamples_test] = 1
index[ind_has_anno] = x
pickle.dump({'index': index}, open(meta_pkl, 'wb'))
np.save(osp.join(data_h5_dir, '../label/label-iter3-train.npy'), label_train)